import random
import math
#from geometry import Vector3	
#from geometry import Quaternion
	
class Vector3(object):
	def __init__(self, vec=None):
		self.v = None
		t = type(vec)
		if t is list or t is tuple:
			self.v = vec[:]
		else:
			self.v = vec.v[:]
			
	def __repr__(self):
		return str(self.v)
	
	def __mul__(self, other):
		v1 = self.v
		t = type(other)
		if t is list or t is tuple:
			v2 = other
		elif t is int or t is float:
			v2 = [other for i in range(3)]
		else:
			v2 = other.v
		return Vector3([vv[0]*vv[1] for vv in zip(v1,v2)])
	
	def __add__(self, other):
		v1 = self.v
		t = type(other)
		if t is list or t is tuple:
			v2 = other
		elif t is int or t is float:
			v2 = [other for i in range(3)]
		else:
			v2 = other.v
		return Vector3([vv[0]+vv[1] for vv in zip(v1,v2)])
		
	def __sub__(self, other):
		v1 = self.v
		t = type(other)
		if t is list or t is tuple:
			v2 = other
		elif t is int or t is float:
			v2 = [other for i in range(3)]
		else:
			v2 = other.v
		return Vector3([vv[0]-vv[1] for vv in zip(v1,v2)])
	
	def dot(self, other):
		v1 = self.v
		t = type(other)
		if t is list or t is tuple:
			v2 = other
		else:
			v2 = other.v
		return sum([vv[0]*vv[1] for vv in zip(v1,v2)])
		
	def cross(self, other):
		v1 = self.v
		t = type(other)
		if t is list or t is tuple:
			v2 = other
		else:
			v2 = other.v
		return Vector3( [v1[1]*v2[2]-v1[2]*v2[1],
						v1[2]*v2[0]-v1[0]*v2[2],
						v1[0]*v2[1]-v1[1]*v2[0]] )
						
	def sq_dist(self, other):
		v1 = self.v
		t = type(other)
		if t is list or t is tuple:
			v2 = other
		else:
			v2 = other.v
		s = [vv[0]-vv[1] for vv in zip(v1,v2)]
		s = [vv*vv for vv in s]
		return sum(s)
						
	def normalize(self):
		s = self.dot(self)
		if s == 1.0:
			return
		s = math.sqrt(s)
		if abs(s-1.0) < 1e-5:
			return
		s = 1. / s
		for i in range(3):
			self.v[i] = self.v[i] * s
	
	
xVector = Vector3([1.,0.,0.])
yVector = Vector3([0.,1.,0.])
zVector = Vector3([0.,0.,1.])
origin = Vector3([0.,0.,0.])

class Quaternion(object):
	def __init__(self, q=None):
		if q and len(q) == 4:
			self.q = q[:]
		else:
			self.q = [0 for i in range(4)]
		
	def __repr__(self):
		return str(self.q)
		
	def normalize(self):
		d = sum([qq*qq for qq in self.q])
		if d == 1.0:
			return
		d = math.sqrt(d)
		if abs(d-1.0) < 1e-5:
			return
		d = 1.0 / d
		for i in range(4):
			self.q[i] = self.q[i]*d
		
	def conj(self):
		w,x,y,z = self.q
		q = Quaternion()
		q.q = [w, -x, -y, -z]
		return q
		
	def set_axis_angle(self, axis, angle):
		w = math.cos(angle * 0.5)
		self.q[0] = w
		sinq = math.sin(angle * 0.5)
		axis = Vector3(axis)
		axis.normalize()
		axis = axis * sinq
		for i in range(3):
			self.q[i+1] = axis.v[i]
		self.normalize()
	
	def axis_angle(self):
		w,x,y,z = self.q
		angle = 2 * math.acos(w)
		d = math.sqrt(1 - (w*w))
		if d == 0:
			axis = Vector3(yVector)
		else:
			d = 1. / d
			axis = Vector3([x*d,y*d,z*d])
		return (axis, angle)
		
	def set_point(self, point):
		self.q[0] = 0
		for i in range(3):
			self.q[i+1] = point.v[i]
			
	def point(self):
		if abs(self.q[0]) >= 1e-4:
			raise RuntimeError("Not a point: " + str(self))
		return Vector3(self.q[1:])
		
	def set_vector_rotation(self, startV, endV):
		startV = Vector3(startV)
		endV = Vector3(endV)
		startV.normalize()
		endV.normalize()
		
		cross = startV.cross(endV)
		dot = startV.dot(endV)
		
		if abs(dot+1.0) <= 1e-4:
			# Rotation of 180 degrees
			# Pick an axis at right angles to startV and angle of 180
			q = Quaternion()
			q.set_vector_rotation(xVector, startV)
			axis = q*yVector
			self.set_axis_angle(axis, math.pi)
		else:
			self.q[0] = 1+dot
			for i in range(3):
				self.q[i+1] = cross.v[i]
			
		self.normalize()

	# A straightforward but slow way of multiplying two quaternions		
#	def __mulold__(self, other):
#		w1 = self.q[0]
#		w2 = other.q[0]
#		v1 = Vector3(self.q[1:])
#		v2 = Vector3(other.q[1:])
#		
#		wr = w1*w2 - v1.dot(v2)
#		vr = v1.cross(v2) + (v2*w1) + (v1*w2)
#		
#		q = Quaternion()
#		q.q[0] = wr
#		for i in range(3):
#			q.q[i+1] = vr.v[i]
#		return q
		
	def __mul__(self, other):
		t = type(other)
		if t is type(self):
			q = Quaternion()
			w1,x1,y1,z1 = self.q
			w2,x2,y2,z2 = other.q
			q.q[0] = w1*w2 - x1*x2 - y1*y2 - z1*z2
			q.q[1] = x1*w2 + w1*x2 + y1*z2 - z1*y2
			q.q[2] = w1*y2 - x1*z2 + y1*w2 + z1*x2
			q.q[3] = w1*z2 + x1*y2 - y1*x2 + z1*w2
			return q
		else:
			w,x,y,z = self.q
			px,py,pz = other.v
			w2 = w*w
			x2 = x*x
			y2 = y*y
			z2 = z*z
			xy = 2*x*y
			xz = 2*x*z
			wx = 2*w*x
			wy = 2*w*y
			wz = 2*w*z
			yz = 2*y*z
			v = Vector3(origin)
			v.v[0] = (w2+x2-y2-z2)*px + (xy-wz)*py + (xz+wy)*pz
			v.v[1] = (xy+wz)*px + (w2-x2+y2-z2)*py + (yz-wx)*pz
			v.v[2] = (xz-wy)*px + (yz+wx)*py + (w2-x2-y2+z2)*pz
			
			return v







####################################
# Robot code
####################################

jointSpace = [(0.0, 2.0*math.pi), (-1.3963, 1.3963), (-2.9671, 2.9671), \
				(0.0, 2.0*math.pi), (-2.9671, 2.9671), (0.0, 2.0*math.pi)]
#jointSpace = [(0.0, 2.0*math.pi), (-1.3963, 1.3963), (-2.9671, 2.9671), \
#				(0.0, 2.0*math.pi), (-1.5708, 1.5708), (0.0, 2.0*math.pi)]
gripSpace = (0.2, 1.5)

mountHeight = 1.0
mountRad = 0.5
armLengths = [2.0,1.5]
armRad = [0.1,0.1]
gripLength = 0.75
gripWidth = 1.5

jointRad = 0.2
gripThick = 0.1

armMax = sum(armLengths)
armMax += gripLength
armMax2 = armMax*armMax

# Returns a list:
#  -Center of gripper arms' ends
#  -End of gripper arm 1
#  -End of gripper arm 2
#  -Base of gripper arm 1
#  -Base of gripper arm 2
#  -Arm 1 side of gripper base
#  -Arm 2 side of gripper base
#  -Center of gripper base (wrist)
# And a Quaternion:
#  -Local rotation of gripper structure
def gripper_positions(joints, grip, testRotation=False):

	positions = [None]*8
	positions[0] = Vector3([0., gripLength, 0.])
	positions[1] = Vector3([0., gripLength, grip/2.0])
	positions[2] = Vector3([0., gripLength, -grip/2.0])
	positions[3] = Vector3([0., 0., grip/2.0])
	positions[4] = Vector3([0., 0., -grip/2.0])
	positions[5] = Vector3([0., 0., gripWidth/2.0])
	positions[6] = Vector3([0., 0., -gripWidth/2.0])
	positions[7] = Vector3(origin)
	
#	if testRotation:
#		sPositions = [Quaternion(p.q) for p in positions]
	
	joint = Quaternion()
	
	# flex, then spin about local y, then spin about global y
	joint.set_axis_angle(zVector, joints[-2])
	for i in range(len(positions)):
		positions[i] = joint*positions[i]
	
	v = joint*yVector
	v.normalize()
	
	joint.set_axis_angle(v, joints[-1])
	for i in range(len(positions)):
		positions[i] = joint*positions[i]
	
	joint.set_axis_angle(yVector, joints[-3])
	for i in range(len(positions)):
		positions[i] = joint*positions[i]
		
	height = Vector3([0, armLengths[1], 0])
	for i in range(len(positions)):
		positions[i] = positions[i] + height
	
	joint.set_axis_angle(zVector, joints[-4])
	for i in range(len(positions)):
		positions[i] = joint*positions[i]
		
	height = Vector3([0, armLengths[0], 0])
	for i in range(len(positions)):
		positions[i] = positions[i] + height
	
	joint.set_axis_angle(zVector, joints[-5])
	for i in range(len(positions)):
		positions[i] = joint*positions[i]
		
	height = Vector3([0, mountHeight, 0])
	for i in range(len(positions)):
		positions[i] = positions[i] + height
	
	joint.set_axis_angle(yVector, joints[-6])
	for i in range(len(positions)):
		positions[i] = joint*positions[i]
		
	gripRotate = Quaternion()
	gripRotate.set_vector_rotation(yVector, positions[1]-positions[3])
	gripRotate2 = Quaternion()
	gripRotate2.set_vector_rotation(gripRotate * zVector, positions[5]-positions[6])
	gripRotate = gripRotate2*gripRotate
	
#	if testRotation:
#		testPositions = [gripRotate*p*gripRotate.conj() for p in sPositions]
#		diff = positions[-1] - testPositions[-1].point()
#		testPositions = [p.point()+diff for p in testPositions]
#		for i,p in enumerate(testPositions):
#			for j in range(3):
#				if abs(p.v[j] - positions[i].v[j]) >= 1e-2:
#					raise RuntimeError("Positions don't match:\n" + \
#						str(testPositions) + "\n\n" + str(positions) + "\n")
	
	return positions, gripRotate







####################################
# Block code
####################################


#Block 0 is the table
numBlocks = 4

stackOffset = 1e-3

stabilityMargin = 1e-2

whd = [(10,0.001,10)] + [(0.5,0.5,0.5)]*numBlocks
whdV = [Vector3(whdi) for whdi in whd]
masses = [0] + [1.0]*numBlocks

initCenters = [(0.0,-0.001,0.0)] + \
					[(3,whd[1][1]+stackOffset+(i*(1.0+stackOffset)),0.0) \
						for i in range(numBlocks-1)]
initRotations = [Quaternion() for i in range(numBlocks)]
for _r in initRotations:
	for _i in range(3):
		_r.set_axis_angle(Vector3([0.,1.,0.]),0.)
del _r
del _i
# aboveBelow[i][0] is the block i is above, [i][1] is the block i is below
initAboveBelow = [[i, i+2] for i in range(-1, numBlocks-1)]
initAboveBelow[-1][1] = -1
initAboveBelow[0][1] = -1
print "initAboveBelow:",initAboveBelow


spaceBounds = [(-10.0,10.0), (0.0,max(20.0, 2.0*sum([d[2] for d in whd]))), \
				(-10.0,10.0)]
				

# Standard goal configuration
goalAboveBelow = [ab[:] for ab in initAboveBelow]
goalAboveBelow[1] = [numBlocks-1, -1]
goalAboveBelow[2] = [0,3]
goalAboveBelow[numBlocks-1] = [numBlocks-2, 1]

# Goals to just set the top block down
#goalAboveBelow[numBlocks-2] = [numBlocks-3, -1]
#goalAboveBelow[numBlocks-1] = [0,-1]

# Goals for the two-block scenario (ie numBlocks = 3)
#goalAboveBelow = [ [-1,-1], [2, -1], [0, 1] ]
print "Goal:", goalAboveBelow


def extents_collide(ext1, ext2):
	return ext1[1] >= ext2[0] and ext1[0] <= ext2[1]
	
def box_extents(axis, center, whd, boxAxes):
	p = center.dot(axis)
	r = sum([abs(boxAxes[i].dot(axis))*whd.v[i] for i in range(3)])
	return (p-r, p+r)
	
def axis_overlap(axis, center1, center2, whd1, whd2, axes1, axes2):
	ext1 = box_extents(axis, center1, whd1, axes1)
	ext2 = box_extents(axis, center2, whd2, axes2)
	
	return extents_collide(ext1,ext2)
	
def colliding_boxes(axes1, axes2, center1, center2, whd1, whd2):
	for i in range(3):
		if not axis_overlap(axes1[i], center1, center2, whd1, whd2, 
					axes1, axes2):
			return False
		if not axis_overlap(axes2[i], center1, center2, whd1, whd2,
					axes1, axes2):
			return False
		for j in range(3):
			if not axis_overlap(axes1[i].cross(axes2[j]), center1, center2,
					whd1, whd2, axes1, axes2):
				return False
	return True

def colliding_block(centers, rotations, block1, block2=None):	
	axes = [xVector, yVector, zVector]
	axes1 = [rotations[block1]*a for a in axes]
	
	if block2 == None:
		for k in range(numBlocks):
			if k == block1:
				continue
			axes2 = [rotations[k]*a for a in axes]
			
			if colliding_boxes(axes1, axes2, Vector3(centers[block1]),
								Vector3(centers[k]), Vector3(whd[block1]),
								Vector3(whd[k])):
				return True
		return False
		
	axes2 = [rotations[block2]*a for a in axes]
	
	return colliding_boxes(axes1, axes2, Vector3(centers[block1]), 
							Vector3(centers[block2]),
							Vector3(whd[block1]), Vector3(whd[block2]))
	
def colliding(centers, rotations, aboveBelow):
	for i in range(len(centers)):
		for j in range(i):
			if i==aboveBelow[j][0] or j==aboveBelow[i][0]:
				continue
			if colliding_block(centers, rotations, i,j):
				return True
	return False
	
def colliding_robot(centers, rotations):
	cornerMul = [ Vector3([1,0,1]), Vector3([-1,0,1]), Vector3([1,0,-1]),
					Vector3([-1,0,-1]) ]
	edgesMul = [ Vector3([1,0,0]), Vector3([-1,0,0]), Vector3([0,0,-1]),
					Vector3([0,0,1]) ]
	for i in range(1, len(centers)):
		dim = Vector3(whd[i])
		corners = [ dim * m for m in cornerMul ]
		edges = [dim * e for e in edgesMul]
		centerXZ = cornerMul[0]*centers[i]
		for e in edges:
			p = rotations[i] * e
			p = p + centerXZ
			if p.dot( p ) <= mountRad*mountRad:
				return True
		for c in corners:
			p = rotations[i] * c
			p = p + centerXZ
			if p.dot( p ) <= mountRad*mountRad:
				return True
	return False
	
def stable(centers, rotations, aboveBelow):
#		print "Checking stability..."
	stacks = []
	stackDict = {}
	for i,ab in enumerate(aboveBelow):
		if ab[0] == 0:
			assert i not in stackDict
			stacks.append([i])
			stackDict[i] = stacks[-1]
			j = ab[1]
			while j != -1:
				stackDict[j] = stacks[-1]
				stacks[-1].append(j)
				j = aboveBelow[j][1]
	for stack in stacks:
		cumulMass = 0.0
		cumulCMx = 0.0
		cumulCMz = 0.0
		for i in reversed(stack):
			cumulMass += masses[i]
			cumulCMx += centers[i][0]
			cumulCMz += centers[i][2]
			cmx = cumulCMx/cumulMass
			cmz = cumulCMz/cumulMass
			j = aboveBelow[i][0]
			
			whdJ = rotations[j] * whdV[j]
			whdI = rotations[i] * whdV[i]
			centerJ = centers[j]
			centerI = centers[i]
			
			basex = [centers[j][0]-abs(whdJ.v[0]), centers[j][0]+abs(whdJ.v[0])]
			basez = [centers[j][2]-abs(whdJ.v[2]), centers[j][2]+abs(whdJ.v[2])]
			basex[0] = max(basex[0],centers[i][0]-abs(whdI.v[0]))
			basex[1] = min(basex[1],centers[i][0]+abs(whdI.v[0]))
			if cmx < basex[0]+stabilityMargin or cmx > basex[1]-stabilityMargin:
				print "Stack unstable:",stack
				print "CM x",cmx
				print "base x:",basex
				return False
			
			basez[0] = max(basez[0],centers[i][2]-abs(whdI.v[2]))
			basez[1] = min(basez[1],centers[i][2]+abs(whdI.v[2]))
			if cmz < basez[0]+stabilityMargin or cmz > basez[1]-stabilityMargin:
				print "Stack unstable:",stack
				print "CM z",cmz
				print "base z:",basez
				return False
	return True
